package com.tianhai.db.index;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import com.tianhai.db.DiskManager;
import com.tianhai.db.Page;

/**
 * B+树索引
 * 支持整数键和页ID值
 * 节点序列化到8KB的Page中
 */
public class BPlusTree {
    
    private final DiskManager diskManager;
    private int rootPageId;
    private final int order; // B+树的阶数
    
    // 元数据页ID（存储根节点页ID等信息）
    private static final int METADATA_PAGE_ID = 0;
    // 数据页起始ID（元数据页之后）
    private static final int DATA_PAGE_START_ID = 1;
    
    /**
     * 创建新的B+树
     * @param diskManager 磁盘管理器
     * @param order B+树的阶数（建议64-128）
     */
    public BPlusTree(DiskManager diskManager, int order) throws IOException {
        this.diskManager = diskManager;
        this.order = order;
        
        // 尝试加载元数据
        if (!loadMetadata()) {
            // 如果元数据不存在，创建新的根节点
            createNewTree();
        }
    }
    
    /**
     * 创建新的B+树
     */
    private void createNewTree() throws IOException {
        // 确保先分配元数据页（页ID 0）
        if (diskManager.getNextPageId() == 0) {
            diskManager.allocateNewPage(); // 分配元数据页
        }
        
        // 分配根节点页（页ID >= 1）
        Page rootPage = diskManager.allocateNewPage();
        rootPageId = rootPage.getPageId();
        
        // 创建根节点（初始为叶子节点）
        BPlusTreeNode rootNode = new BPlusTreeNode(rootPageId, BPlusTreeNode.NodeType.LEAF, order);
        rootNode.setRoot(true);
        
        // 保存根节点
        saveNode(rootNode);
        
        // 保存元数据
        saveMetadata();
    }
    
    /**
     * 加载元数据
     */
    private boolean loadMetadata() throws IOException {
        try {
            Page metadataPage = diskManager.readPage(METADATA_PAGE_ID);
            byte[] data = metadataPage.getData();
            
            // 读取根节点页ID（从开头读取）
            int offset = 0;
            rootPageId = ((data[offset] & 0xFF) << 24) |
                        ((data[offset + 1] & 0xFF) << 16) |
                        ((data[offset + 2] & 0xFF) << 8) |
                        (data[offset + 3] & 0xFF);
            
            return rootPageId >= DATA_PAGE_START_ID;
        } catch (Exception e) {
            return false;
        }
    }
    
    /**
     * 保存元数据
     */
    private void saveMetadata() throws IOException {
        Page metadataPage;
        try {
            metadataPage = diskManager.readPage(METADATA_PAGE_ID);
        } catch (Exception e) {
            // 创建元数据页（确保是页ID 0）
            metadataPage = new Page(METADATA_PAGE_ID);
        }
        
        byte[] data = metadataPage.getData();
        
        // 写入根节点页ID（从开头写入）
        int offset = 0;
        data[offset] = (byte) (rootPageId >> 24);
        data[offset + 1] = (byte) (rootPageId >> 16);
        data[offset + 2] = (byte) (rootPageId >> 8);
        data[offset + 3] = (byte) rootPageId;
        
        // 写入B+树阶数（用于验证）
        data[4] = (byte) (order >> 24);
        data[5] = (byte) (order >> 16);
        data[6] = (byte) (order >> 8);
        data[7] = (byte) order;
        
        diskManager.writePage(metadataPage);
    }
    
    /**
     * 插入键值对
     * @param key 键（整数）
     * @param value 值（页ID）
     */
    public void insert(int key, int value) throws IOException {
        BPlusTreeNode root = loadNode(rootPageId);
        
        // 检查根节点是否需要分裂
        if (root.isFull()) {
            // 创建新的根节点
            Page newRootPage = diskManager.allocateNewPage();
            int newRootPageId = newRootPage.getPageId();
            
            BPlusTreeNode newRoot = new BPlusTreeNode(
                newRootPageId, 
                BPlusTreeNode.NodeType.INTERNAL, 
                order
            );
            newRoot.setRoot(true);
            // 重要：内部节点的values数量比keys多1，初始时有一个指向旧根的指针
            newRoot.getValues().add(rootPageId);
            
            // 旧根节点不再是根
            root.setRoot(false);
            // 注意：不在这里保存root，因为splitChild会修改并保存root
            
            // 分裂旧根节点（此时newRoot在内存中会被修改）
            splitChild(newRoot, 0, root);
            
            // 更新根节点ID并保存元数据
            rootPageId = newRootPageId;
            saveMetadata();
            
            // 注意：newRoot已经在splitChild中被修改和保存，不需要重新加载
            // 直接在当前newRoot对象中插入
            insertNonFull(newRoot, key, value);
        } else {
            insertNonFull(root, key, value);
        }
    }
    
    /**
     * 在非满节点中插入
     */
    private void insertNonFull(BPlusTreeNode node, int key, int value) throws IOException {
        if (node.getNodeType() == BPlusTreeNode.NodeType.LEAF) {
            // 叶子节点：直接插入
            node.insertInLeaf(key, value);
            saveNode(node);
        } else {
            // 内部节点：找到合适的子节点
            int pos = 0;
            while (pos < node.getKeys().size() && key >= node.getKeys().get(pos)) {
                pos++;
            }
            
            // 确保 pos 不超出values的范围
            if (node.getValues().isEmpty()) {
                throw new IllegalStateException(
                    "内部节点的values列表为空: pageId=" + node.getPageId() + 
                    ", keys.size=" + node.getKeys().size() + 
                    ", values.size=" + node.getValues().size()
                );
            }
            if (pos >= node.getValues().size()) {
                pos = node.getValues().size() - 1;
            }
            if (pos < 0) {
                pos = 0;
            }
            
            int childPageId = node.getValues().get(pos);
            BPlusTreeNode child = loadNode(childPageId);
            
            if (child.isFull()) {
                // 子节点已满，需要分裂
                splitChild(node, pos, child);
                
                // 分裂后，重新确定插入位置
                // 注意：node已经被修改，不需要重新加载
                if (pos < node.getKeys().size() && key >= node.getKeys().get(pos)) {
                    pos++;
                }
                if (pos >= node.getValues().size()) {
                    pos = node.getValues().size() - 1;
                }
                if (pos < 0) {
                    pos = 0;
                }
                childPageId = node.getValues().get(pos);
                child = loadNode(childPageId);
            }
            
            insertNonFull(child, key, value);
        }
    }
    
    /**
     * 分裂子节点
     */
    private void splitChild(BPlusTreeNode parent, int childIndex, BPlusTreeNode child) 
            throws IOException {
        // 分配新页
        Page newPage = diskManager.allocateNewPage();
        int newPageId = newPage.getPageId();
        
        // 分裂子节点
        BPlusTreeNode.SplitResult result = child.split(newPageId);
        
        // 先更新父节点（在内存中）
        parent.insertInInternal(result.splitKey, newPageId);
        
        // 然后保存所有节点
        saveNode(child);
        saveNode(result.newNode);
        saveNode(parent);
    }
    
    /**
     * 搜索键对应的值
     * @param key 要搜索的键
     * @return 对应的值（页ID），如果不存在返回null
     */
    public Integer search(int key) throws IOException {
        return searchInNode(rootPageId, key);
    }
    
    /**
     * 在节点中搜索
     */
    private Integer searchInNode(int nodePageId, int key) throws IOException {
        BPlusTreeNode node = loadNode(nodePageId);
        
        if (node.getNodeType() == BPlusTreeNode.NodeType.LEAF) {
            // 叶子节点：查找键
            for (int i = 0; i < node.getKeys().size(); i++) {
                if (node.getKeys().get(i) == key) {
                    return node.getValues().get(i);
                }
            }
            return null;
        } else {
            // 内部节点：找到合适的子节点
            int pos = 0;
            while (pos < node.getKeys().size() && key >= node.getKeys().get(pos)) {
                pos++;
            }
            
            // 内部节点的values数量比keys多1
            if (node.getValues().isEmpty()) {
                throw new IllegalStateException(
                    "内部节点values为空: pageId=" + nodePageId + 
                    ", keys.size=" + node.getKeys().size()
                );
            }
            if (pos < 0) pos = 0;
            if (pos >= node.getValues().size()) pos = node.getValues().size() - 1;
            
            int childPageId = node.getValues().get(pos);
            return searchInNode(childPageId, key);
        }
    }
    
    /**
     * 范围搜索
     * @param startKey 起始键（包含）
     * @param endKey 结束键（包含）
     * @return 范围内的所有键值对
     */
    public List<KeyValuePair> rangeSearch(int startKey, int endKey) throws IOException {
        List<KeyValuePair> results = new ArrayList<>();
        
        // 找到第一个叶子节点
        BPlusTreeNode leafNode = findLeafNode(rootPageId, startKey);
        
        // 遍历叶子节点链表
        while (leafNode != null) {
            for (int i = 0; i < leafNode.getKeys().size(); i++) {
                int key = leafNode.getKeys().get(i);
                
                if (key > endKey) {
                    return results;
                }
                
                if (key >= startKey) {
                    results.add(new KeyValuePair(key, leafNode.getValues().get(i)));
                }
            }
            
            // 移动到下一个叶子节点
            if (leafNode.getNextLeafPageId() != -1) {
                leafNode = loadNode(leafNode.getNextLeafPageId());
            } else {
                break;
            }
        }
        
        return results;
    }
    
    /**
     * 找到包含指定键的叶子节点
     */
    private BPlusTreeNode findLeafNode(int nodePageId, int key) throws IOException {
        BPlusTreeNode node = loadNode(nodePageId);
        
        if (node.getNodeType() == BPlusTreeNode.NodeType.LEAF) {
            return node;
        } else {
            int pos = 0;
            while (pos < node.getKeys().size() && key >= node.getKeys().get(pos)) {
                pos++;
            }
            
            // 内部节点的values数量比keys多1
            if (pos < 0 || pos >= node.getValues().size()) {
                if (pos < 0) pos = 0;
                if (pos >= node.getValues().size()) pos = node.getValues().size() - 1;
            }
            
            int childPageId = node.getValues().get(pos);
            return findLeafNode(childPageId, key);
        }
    }
    
    /**
     * 从磁盘加载节点
     */
    private BPlusTreeNode loadNode(int pageId) throws IOException {
        // 使用rawReadPage直接读取原始字节，不经过Page的头部解析
        byte[] data = diskManager.rawReadPage(pageId);
        return BPlusTreeNode.deserialize(data, pageId, order);
    }
    
    /**
     * 保存节点到磁盘
     */
    private void saveNode(BPlusTreeNode node) throws IOException {
        byte[] nodeData = node.serialize();
        
        if (nodeData.length != Page.PAGE_SIZE) {
            throw new IllegalStateException(
                "B+树节点序列化后的数据大小 (" + nodeData.length + 
                ") 与页大小 (" + Page.PAGE_SIZE + ") 不匹配"
            );
        }
        
        // 使用不解析头部的构造函数，parseHeader=false
        Page page = new Page(node.getPageId(), nodeData, false);
        
        diskManager.writePage(page);
    }
    
    /**
     * 获取根节点页ID
     */
    public int getRootPageId() {
        return rootPageId;
    }
    
    /**
     * 获取B+树的阶数
     */
    public int getOrder() {
        return order;
    }
    
    /**
     * 键值对结果
     */
    public static class KeyValuePair {
        public final int key;
        public final int value;
        
        public KeyValuePair(int key, int value) {
            this.key = key;
            this.value = value;
        }
        
        @Override
        public String toString() {
            return "(" + key + ", " + value + ")";
        }
    }
}
